#include "dwa.h"

DWAPlanner::DWAPlanner()
{
    ros::NodeHandle nh_;
    trajectories_viz_pub = nh_.advertise<visualization_msgs::MarkerArray>("/dwa_trajs", 10);
    best_traj_viz_pub = nh_.advertise<visualization_msgs::MarkerArray>("/best_traj", 10);
    cmd_vel_pub_ = nh_.advertise<geometry_msgs::Twist>("/cmd_vel", 10);
    path_viz_pub = nh_.advertise<visualization_msgs::Marker>("/path_viz", 10);

    // change var names
    nh_.param("HZ", HZ, 20.0);
    nh_.param("TARGET_VELOCITY", TARGET_VELOCITY, 0.8);
    nh_.param("MAX_VELOCITY", MAX_VELOCITY, 1.0);
    nh_.param("MIN_VELOCITY", MIN_VELOCITY, 0.0);
    nh_.param("MAX_OMEGA", MAX_OMEGA, 0.8);
    nh_.param("MAX_ACCELERATION", MAX_ACCELERATION, 2.0);
    nh_.param("MAX_ANG_ACCELERATION", MAX_ANG_ACCELERATION, 3.0);
    //local_nh.param("MAX_DIST", MAX_DIST, {10.0});
    nh_.param("VELOCITY_RES", VELOCITY_RES, 0.05);
    nh_.param("OMEGA_RES", OMEGA_RES, 0.05);
    nh_.param("SIM_TIME", SIM_TIME, 3.0);
    nh_.param("TO_GOAL_COST_GAIN", TO_GOAL_COST_GAIN, 1.0);
    nh_.param("SPEED_COST_GAIN", SPEED_COST_GAIN, 1.0);
    nh_.param("OBSTACLE_COST_GAIN", OBSTACLE_COST_GAIN, 1.0);
    nh_.param("GOAL_THRESHOLD", GOAL_THRESHOLD, 0.3);
    DT = 1.0 / HZ;
}

Window DWAPlanner::get_window(double curr_vel, double curr_w)
{
    //
    curr_vel_.linear.x = curr_vel;
    curr_vel_.angular.z = curr_w;

    //
    Window window(MIN_VELOCITY, MAX_VELOCITY, -MAX_OMEGA, MAX_OMEGA);
    // edit presentation
    window.min_vel = std::max((curr_vel - MAX_ACCELERATION * DT), MIN_VELOCITY);
    window.max_vel = std::min((curr_vel + MAX_ACCELERATION * DT), MAX_VELOCITY);
    window.min_omega = std::max((curr_w - MAX_ANG_ACCELERATION * DT), -MAX_OMEGA);
    window.max_omega = std::min((curr_w + MAX_ANG_ACCELERATION * DT), MAX_OMEGA);

    return window;
}

std::vector<State> DWAPlanner::best_dwa_selection(double x_init, double y_init, double theta_init, const Window &window, const Eigen::Vector3d &goal,
                                                  const std::vector<std::vector<double>> &local_obs)
{

    dwa_converged_ = false;
    double min_cost = DBL_MAX;

    std::vector<std::vector<State>> trajectory_table;
    std::vector<State> best_trajectory;

    // creating candidate trajectories
    for (double v = window.min_vel; v <= window.max_vel; v += VELOCITY_RES)
    {
        for (double w = window.min_omega; w <= window.max_omega; w += OMEGA_RES)
        {

            State state(x_init, y_init, theta_init, curr_vel_.linear.x, curr_vel_.angular.z);
            std::vector<State> temp_traj;
            for (double t = 0.0; t <= SIM_TIME; t += DT)
            {
                simulate_dynamics(state, v, w);
                temp_traj.push_back(state);
            }
            trajectory_table.push_back(temp_traj);

            // edit var names and method names
            double cost_on_goal = get_goal_cost(temp_traj, goal);
            double cost_on_speed = get_speed_cost(temp_traj, TARGET_VELOCITY);
            double cost_on_obs = get_obstacle_cost(temp_traj, local_obs);

            double cumulative_cost = TO_GOAL_COST_GAIN * cost_on_goal + SPEED_COST_GAIN * cost_on_speed + OBSTACLE_COST_GAIN * cost_on_obs;

            if (cumulative_cost <= min_cost)
            {
                min_cost = cumulative_cost;
                best_trajectory = temp_traj;
                dwa_converged_ = true;
            }
        }
    }

    //trajectory table populated
    show_dwa_trajectories(trajectory_table);
    // ROS_INFO("Traj Table Created and Published");

    // updating minimum cost trajectory
    if (min_cost == DBL_MAX)
    {
        std::vector<State> traj;
        State state(x_init, y_init, theta_init, curr_vel_.linear.x, curr_vel_.angular.z);
        traj.push_back(state);
        dwa_converged_ = false;
        ROS_ERROR("INFINITE COST");
        best_trajectory = traj;
    }

    return best_trajectory;
}

void DWAPlanner::show_dwa_trajectories(const std::vector<std::vector<State>> &trajectory_table)
{
    visualization_msgs::MarkerArray traj_list;
    visualization_msgs::Marker traj;
    geometry_msgs::Point p;

    traj.header.frame_id = Frame_id;
    traj.header.stamp = ros::Time::now();
    traj.id = 0;
    traj.type = visualization_msgs::Marker::LINE_STRIP;

    traj.pose.position.x = 0.0;
    traj.pose.position.y = 0.0;
    traj.pose.position.z = 0;
    traj.pose.orientation.x = 0.0;
    traj.pose.orientation.y = 0.0;
    traj.pose.orientation.z = 0.0;
    traj.pose.orientation.w = 1.0;

    traj.scale.x = traj.scale.y = traj.scale.z = 0.01;
    traj.color.r = 1.0;
    traj.color.a = 0.5;

    for (int i = 0; i < trajectory_table.size(); i++)
    {

        traj.id += 1;
        traj.color.b += 0.1;
        traj.points.clear();

        for (int j = 0; j < trajectory_table[i].size(); j++)
        {
            p.x = trajectory_table[i][j].x;
            p.y = trajectory_table[i][j].y;
            traj.points.push_back(p);
        }
        traj_list.markers.push_back(traj);
    }
    trajectories_viz_pub.publish(traj_list);
}

void DWAPlanner::show_best_trajectory(const std::vector<State> &best_trajectory)
{

    visualization_msgs::MarkerArray best_list;
    visualization_msgs::Marker best_traj;
    geometry_msgs::Point point;

    best_traj.header.frame_id = Frame_id;
    best_traj.id = 1;
    best_traj.type = visualization_msgs::Marker::LINE_STRIP;
    best_traj.scale.x = best_traj.scale.y = 0.02;
    best_traj.action = visualization_msgs::Marker::ADD;
    best_traj.pose.orientation.w = 1.0;
    best_traj.color.g = 1.0;
    best_traj.color.a = 1.0;

    for (int i = 0; i < best_trajectory.size(); i++)
    {
        point.x = best_trajectory[i].x;
        point.y = best_trajectory[i].y;
        best_traj.points.push_back(point);
    }
    best_list.markers.push_back(best_traj);
    best_traj_viz_pub.publish(best_list);
}

void DWAPlanner::simulate_dynamics(State &state, double velocity, double omega)
{
    state.heading += omega * DT;
    state.x += velocity * cos(state.heading) * DT;
    state.y += velocity * sin(state.heading) * DT;
    state.velocity = velocity;
    state.omega = omega;
}

double DWAPlanner::get_goal_cost(const std::vector<State> &current_trajectory, const Eigen::Vector3d &goal)
{

    Eigen::Vector2d traj_end(current_trajectory.back().x, current_trajectory.back().y);
    Eigen::Vector2d goal_res(0.0, 0.0);

    goal_res[0] = traj_end[0] - goal[0];
    goal_res[1] = traj_end[1] - goal[1];

    return goal_res.norm();
}

double DWAPlanner::get_speed_cost(const std::vector<State> &current_trajectory, const double target_vel)
{

    double cost = fabs(target_vel - current_trajectory.end()->velocity);

    return cost;
}

double DWAPlanner::get_obstacle_cost(const std::vector<State> &trajectory, const std::vector<std::vector<double>> &local_obs)
{
    double cost = 0.0;
    double min_dist = DBL_MAX;

    for (const auto &state : trajectory)
    {
        for (const auto &scan_pt : local_obs)
        {
            double current_distance = sqrt((state.x - scan_pt[0]) * (state.x - scan_pt[0]) +
                                           (state.y - scan_pt[1]) * (state.y - scan_pt[1]));
            if (current_distance <= 0.01)
            {

                cost = DBL_MAX;
                return cost;
            }
            min_dist = std::min(min_dist, current_distance);
        }
    }

    if (min_dist == 0.0)
        return DBL_MAX;

    cost = 1.0 / min_dist;
    return cost;
}

DiscretizedTrajectory DWAPlanner::trajectoryconvert(std::vector<State> trajectory_opt)
{
    DiscretizedTrajectory best_path;

    // 将最优的轨迹输出
    for (size_t m = 0; m < trajectory_opt.size(); m++)
    {
        TrajectoryPoint local_lane;
        local_lane.set_x(trajectory_opt[m].x); //x
        local_lane.set_y(trajectory_opt[m].y); // y
        local_lane.set_z(0);                    //z

        local_lane.set_theta(trajectory_opt[m].heading); // theta
        local_lane.set_v(trajectory_opt[m].velocity);     //v

        best_path.push_back(local_lane);
    }
    return best_path;
}